from PyQt5 import QtWidgets
from PyQt5.QtWidgets import QMenu, QAction, QMessageBox
from qgis.core import QgsProject
from qgis.gui import QgsLayerTreeViewMenuProvider, QgsLayerTreeView, QgsLayerTreeViewDefaultActions, QgsMapCanvas

PROJECT = QgsProject.instance()
class menuProvider(QgsLayerTreeViewMenuProvider):
    def __init__(self,mainWindow, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.layerTreeView: QgsLayerTreeView = mainWindow.layerTreeView
        self.mapCanvas: QgsMapCanvas = mainWindow.mapCanvas
        self.mainWindows = mainWindow

    def createContextMenu(self) -> QtWidgets.QMenu:
        menu = QMenu()
        self.actions : QgsLayerTreeViewDefaultActions = self.layerTreeView.defaultActions()
        if self.layerTreeView.currentIndex().isValid():
            menu.addAction('展开所有图层',self.layerTreeView.expandAllNodes)
            menu.addAction('折叠所有图层',self.layerTreeView.collapseAllNodes)
            return menu

        if len(self.layerTreeView.selectedLayers()) > 1:
            # 添加组
            self.actionGroupSelected = self.actions.actionGroupSelected()
            menu.addAction(self.actionGroupSelected)

            actionDeleteSelectedLayers = QAction('删除选中图层',menu)
            actionDeleteSelectedLayers.triggered.connect(self.deleteSelectedLayer)
            menu.addAction(actionDeleteSelectedLayers)

            return menu

    def deleteSelectedLayer(self):
        deleteRes = QMessageBox.question(self.mainWindows, '信息', "确定要删除所选图层？", QMessageBox.Yes | QMessageBox.No,
                                         QMessageBox.No)
        if deleteRes == QMessageBox.Yes:
            layers = self.layerTreeView.selectedLayers()
            for layer in layers:
                self.deleteLayer(layer)

    def deleteLayer(self,layer):
        PROJECT.removeMapLayer(layer)
        self.mapCanvas.refresh()
        return 0